{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data preparation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "from Bio.PDB import *\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from torch_cluster import knn\n",
    "from antiberty import AntiBERTyRunner\n",
    "from torch_geometric.data import Data\n",
    "\n",
    "import sys \n",
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")\n",
    "sys.path.append(r\"E:\\608\\论文\\代码\\Nanotope\\Nanotope\")\n",
    "\n",
    "pdbfile='6h72.pdb'\n",
    "chainid='C'\n",
    "NB_MAX_LENGTH = 140\n",
    "k=32\n",
    "BACKBONE_ATOMS = [\"N\", \"CA\", \"C\", \"O\", \"CB\"]\n",
    "OUTPUT_SIZE = len(BACKBONE_ATOMS) * 3\n",
    "\n",
    "# get seq, residues by chain id from Nanobody PDB file\n",
    "def get_seq_aa(pdb_file, chain_id):\n",
    "    \n",
    "    # get chain\n",
    "    chain = PDBParser(QUIET=True).get_structure(pdb_file, pdb_file)[0][chain_id]\n",
    "    aa_residues = []\n",
    "    seq = \"\"\n",
    "\n",
    "    for residue in chain.get_residues():\n",
    "        aa = residue.get_resname()\n",
    "        if not is_aa(aa) or not residue.has_id('CA'):\n",
    "            continue\n",
    "        elif aa == \"UNK\":  \n",
    "            seq += \"X\"\n",
    "        else:\n",
    "            seq += Polypeptide.three_to_one(residue.get_resname())\n",
    "        aa_residues.append(residue)\n",
    "\n",
    "    return seq, aa_residues\n",
    "\n",
    "# get Nanobody chain C-a atom coordinates\n",
    "def generate_coord(pdb_file,chain_id):  \n",
    "    seq, aa_residues = get_seq_aa(pdb_file, chain_id)\n",
    "    xyz_matrix = np.zeros((NB_MAX_LENGTH, OUTPUT_SIZE))\n",
    "    for i in range(len(aa_residues)):\n",
    "        for j, atom in enumerate(BACKBONE_ATOMS):\n",
    "            if not (atom==\"CB\" and seq[i] == \"G\"):\n",
    "                xyz_matrix[i][3*j:3*j+3] = aa_residues[i][atom].get_coord()\n",
    "\n",
    "    return seq,xyz_matrix[:,3:6] #C-a\n",
    "\n",
    "seq,coord = generate_coord(pdbfile,chainid)\n",
    "size = len(seq)\n",
    "coord = torch.tensor(coord)\n",
    "# get seq embedding by Antiberty model \n",
    "Antiberty = AntiBERTyRunner()\n",
    "embeddings = Antiberty.embed([seq])[0][1:-1]\n",
    "# padding if len(seq)<140, using zero vetor [0,...,0]\n",
    "if size<140:\n",
    "    pad = torch.zeros((140-size),512).cuda()\n",
    "    embeddings =torch.cat([embeddings,pad],dim=0)\n",
    "\n",
    "# construct KNN edges\n",
    "edge_index = knn(coord,coord,k = k)\n",
    "# construct graph data\n",
    "data = Data(x = embeddings.unsqueeze(0),edge_index=edge_index,mask=size,batch=torch.Tensor([0]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## model prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.GNNnet import Nanotope\n",
    "\n",
    "model = Nanotope(hidden_channels=512, num_layers=3, num_heads=8,num_bases=8)\n",
    "model.load_state_dict(torch.load('../model/model_weights/model_weights.pt'))\n",
    "model.eval().cuda()\n",
    "\n",
    "#prediction\n",
    "prediction = model(data)[0:data.mask]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Q \t 0.543\n",
      "V \t 0.753\n",
      "Q \t 0.632\n",
      "L \t 0.629\n",
      "Q \t 0.526\n",
      "E \t 0.644\n",
      "S \t 0.577\n",
      "G \t 0.554\n",
      "G \t 0.657\n",
      "G \t 0.528\n",
      "L \t 0.388\n",
      "V \t 0.339\n",
      "Q \t 0.246\n",
      "A \t 0.295\n",
      "G \t 0.275\n",
      "G \t 0.447\n",
      "S \t 0.386\n",
      "L \t 0.447\n",
      "R \t 0.362\n",
      "L \t 0.469\n",
      "S \t 0.485\n",
      "C \t 0.595\n",
      "A \t 0.206\n",
      "A \t 0.550\n",
      "S \t 0.042\n",
      "G \t 0.191\n",
      "R \t 0.366\n",
      "M \t 0.502\n",
      "F \t 0.295\n",
      "S \t 0.506\n",
      "I \t 0.812\n",
      "N \t 0.603\n",
      "S \t 0.601\n",
      "M \t 0.756\n",
      "G \t 0.851\n",
      "W \t 0.941\n",
      "Y \t 0.776\n",
      "R \t 0.590\n",
      "Q \t 0.415\n",
      "A \t 0.309\n",
      "P \t 0.188\n",
      "G \t 0.120\n",
      "K \t 0.127\n",
      "E \t 0.237\n",
      "R \t 0.536\n",
      "E \t 0.314\n",
      "L \t 0.791\n",
      "V \t 0.589\n",
      "A \t 0.784\n",
      "T \t 0.920\n",
      "I \t 0.609\n",
      "S \t 0.852\n",
      "E \t 0.798\n",
      "A \t 0.536\n",
      "G \t 0.308\n",
      "T \t 0.773\n",
      "T \t 0.485\n",
      "T \t 0.836\n",
      "Y \t 0.683\n",
      "A \t 0.484\n",
      "D \t 0.396\n",
      "S \t 0.338\n",
      "V \t 0.524\n",
      "R \t 0.632\n",
      "G \t 0.348\n",
      "R \t 0.440\n",
      "F \t 0.516\n",
      "T \t 0.479\n",
      "I \t 0.642\n",
      "A \t 0.506\n",
      "R \t 0.730\n",
      "D \t 0.090\n",
      "N \t 0.204\n",
      "A \t 0.163\n",
      "K \t 0.111\n",
      "N \t 0.094\n",
      "T \t 0.324\n",
      "V \t 0.736\n",
      "Y \t 0.744\n",
      "L \t 0.777\n",
      "Q \t 0.533\n",
      "M \t 0.762\n",
      "N \t 0.548\n",
      "S \t 0.374\n",
      "L \t 0.445\n",
      "N \t 0.435\n",
      "P \t 0.261\n",
      "E \t 0.407\n",
      "D \t 0.667\n",
      "T \t 0.457\n",
      "A \t 0.405\n",
      "V \t 0.420\n",
      "Y \t 0.440\n",
      "Y \t 0.543\n",
      "C \t 0.420\n",
      "N \t 0.775\n",
      "A \t 0.502\n",
      "Y \t 0.811\n",
      "I \t 0.602\n",
      "Q \t 0.691\n",
      "L \t 0.833\n",
      "D \t 0.710\n",
      "S \t 0.686\n",
      "T \t 0.647\n",
      "I \t 0.466\n",
      "W \t 0.708\n",
      "F \t 0.557\n",
      "R \t 0.578\n",
      "A \t 0.574\n",
      "Y \t 0.401\n",
      "W \t 0.604\n",
      "G \t 0.496\n",
      "Q \t 0.478\n",
      "G \t 0.479\n",
      "T \t 0.533\n",
      "Q \t 0.410\n",
      "V \t 0.443\n",
      "T \t 0.377\n",
      "V \t 0.418\n",
      "S \t 0.268\n",
      "S \t 0.172\n"
     ]
    }
   ],
   "source": [
    "def display_prediction(seq,prediction):\n",
    "    for a,pre in zip(seq,prediction):\n",
    "        print(a,'\\t','%.3f'%pre.item())\n",
    "\n",
    "display_prediction(seq,prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "ec3b6344b8b3f6c3db1463248c46b8920f4fc68f6187e2c1e71f56a52299e245"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
